import os
import threading
import time

import allure
import pyautogui
from PIL import Image
from selenium import webdriver

from common.log import log

Flag = False


class RecordGif:
    """
    截图并制作Gif动画
    """

    @allure.step('创建或清空目录')
    def clear_dir(self, path):
        # 创建或清空目录
        if not os.path.isdir(path):
            os.makedirs(path)  # 创建目录
        else:  # 清空目录
            [os.remove(os.path.join(path, file_name)) for file_name in os.listdir(path)]

    @allure.step('web的方式截图并制作Gif动画')
    def web_record_shot(self, dr, img_dir):
        """
        web的方式截图并制作Gif动画
        循环截图函数
        :param dr:
        :param img_dir:
        :return:
        """
        i = 0
        self.clear_dir(img_dir)  # 清空目录
        while True:
            i += 1
            time.sleep(1)
            img_file = os.path.join(img_dir, f'{i}.png')
            try:
                dr.save_screenshot(img_file)
            except:
                return

    @allure.step('第三方库 的方式截图并制作Gif动画')
    def shot(self, img_dir):
        """
        第三方库 的方式截图并制作Gif动画
        循环截图函数
        :param img_dir:
        :return:
        """
        i = 0
        self.clear_dir(img_dir)  # 清空目录
        time.sleep(3)
        while True:
            i += 1
            img_file = os.path.join(img_dir, f'{i}.jpg')
            time.sleep(1)
            im = pyautogui.screenshot()
            im.save(img_file)
            if flag:
                log.info("截图结束！")
                break

    @allure.step('停止截图')
    def stop_shot(key=0):
        global flag
        if key == 1:
            flag = True
            log.info("停止截图！")
            return False


if __name__ == '__main__':
    # Selenium操作
    img_dir = '../img'  # 临时图片目录
    dr = webdriver.Chrome()
    t = threading.Thread(target=RecordGif().web_record_shot, args=(dr, img_dir), )  # 新建线程
    t.start()  # 启动截图线程
    dr.get('https://www.baidu.com')

    list_key = ['test', 'hh', 'weo', 'woe']
    for i in list_key:
        dr.find_element("id", "kw").send_keys(i)
        dr.find_element('id', 'su').click()
        time.sleep(1)
    time.sleep(1)
    dr.back()
    time.sleep(2)
    dr.quit()

    # 图片拼接成gif
    img_list = os.listdir(img_dir)  # 列出目录所有图片
    img_list.sort(key=lambda x: int(x[:-4]))  # 排序
    first_img = Image.open(os.path.join(img_dir, img_list[0]))  # 第一张图片对象
    else_imgs = [Image.open(os.path.join(img_dir, img)) for img in img_list[1:]]  # 剩余图片对象
    first_img.save("record.gif", append_images=else_imgs,
                   duration=300,  # 每张图片的过过渡时间
                   save_all=True)  # 拼接保存，如果想要循环播放可以加上loop=0

    video_file = 'D:\\worktest\\record.gif'
    with open(video_file, "rb") as f:
        allure.attach(f.read(), name="GIF", attachment_type=allure.attachment_type.GIF)
